"""
This is part of Human Robot Dynamics-Control, By Tourists.

Based on "OpenLoong Dynamics Control"
Transfer from mujoco to ISAAC SIM by IsaacLab

"""
from ..common import DataBus
from ..common import LegState
from ..common import MotionState
import numpy as np 
from scipy.linalg import pinv

#pragma once

#include "data_bus.h"
#include <Eigen/Dense>
#include "useful_math.h"

#pragma once

#include "data_bus.h"
#include <Eigen/Dense>
#include "useful_math.h"

class GaitScheduler:
    def __init__(self, tSwingIn, dtIn):
        self.tSwing = tSwingIn
        self.dt = dtIn
        self.isIni = False

        self.phi = 0.0
        # self.tSwing = 0.4
        # self.dt = 0.001
        self.FzThrehold = 100
        self.Fz_L_m = 0
        self.Fz_R_m = 0
        self.legState = LegState.RSt
        self.legStateNext = LegState.RSt
        self.motionState = MotionState.Stand
        # GaitScheduler(double tSwingIn, double dtIn);
        # void dataBusRead(const DataBus &robotState);
        # void dataBusWrite(DataBus &robotState);
        # void step();
        # void stop();
        self.FLest = np.array([])
        self.FRest = np.array([])
        self.torJoint = np.array([])

        self.enableNextStep = False
        self.touchDown = False ## touch down event indicator
        
        # private:
        self.fe_r_pos_W = np.array([]) 
        self.fe_l_pos_W = np.array([]) 
        self.swingStartPos_W = np.array([]) 
        self.posHip_W = np.array([]) 
        self.posST_W = np.array([]) 
        self.hip_r_pos_W = np.array([]) 
        self.hip_l_pos_W = np.array([]) 
        self.dq = np.array([]) 
        self.stanceStartPos_W = np.array([])
        self.fe_r_rot_W = np.array([]) 
        self.fe_l_rot_W = np.array([])
        self.dyn_M = np.array([]) 
        self.dyn_Non = np.array([]) 
        self.J_l = np.array([]) 
        self.J_r = np.array([]) 
        self.dJ_l = np.array([]) 
        self.dJ_r = np.array([])
        self.theta0 = 0.0 
        self.model_nv = 0

    ## Note: no double-support here, swing time always equals to stance time


    def dataBusRead(self, robotState):
        self.model_nv = robotState.model_nv
        self.torJoint = np.zeros(self.model_nv - 6)
        for i in range(self.model_nv - 6):
            self.torJoint[i] = robotState.motors_tor_cur[i]
        self.dyn_M = robotState.dyn_M
        self.dyn_Non = robotState.dyn_Non
        self.J_l = robotState.J_l
        self.dJ_l = robotState.dJ_l
        self.J_r = robotState.J_r
        self.dJ_r = robotState.dJ_r
        self.Fz_L_m = robotState.fL[2]
        self.Fz_R_m = robotState.fR[2]
        self.hip_l_pos_W = robotState.hip_l_pos_W
        self.hip_r_pos_W = robotState.hip_r_pos_W
        self.fe_r_pos_W = robotState.fe_r_pos_W
        self.fe_l_pos_W = robotState.fe_l_pos_W
        self.fe_l_rot_W = robotState.fe_l_rot_W
        self.fe_r_rot_W = robotState.fe_r_rot_W
        self.dq = robotState.dq
        self.motionState = robotState.motionState

    def dataBusWrite(self, robotState):
        robotState.tSwing = self.tSwing
        robotState.swingStartPos_W = self.swingStartPos_W
        robotState.stanceDesPos_W = self.stanceStartPos_W
        robotState.posHip_W = self.posHip_W
        robotState.posST_W = self.posST_W
        robotState.theta0 = self.theta0
        robotState.legState = self.legState
        robotState.legStateNext = self.legStateNext
        robotState.phi = self.phi
        robotState.FL_est = self.FLest
        robotState.FR_est = self.FRest
        if self.legState == LegState.LSt:
            robotState.stance_fe_pos_cur_W = self.fe_l_pos_W
            robotState.stance_fe_rot_cur_W = self.fe_l_rot_W
        else:
            robotState.stance_fe_pos_cur_W = self.fe_r_pos_W
            robotState.stance_fe_rot_cur_W = self.fe_r_rot_W
        robotState.motionState = self.motionState

    def step(self):
        tauAll = np.zeros(self.model_nv)
        tauAll[6:] = self.torJoint
        # print("self.J_l shape:", self.J_l.shape)
        # print("self.dyn_M shape:", self.dyn_M.shape)
        # print("self.dJ_l shape:", self.dJ_l.shape)
        # print("self.dq shape:", self.dq.shape)
        self.FLest = -pinv(self.J_l @ np.linalg.inv(self.dyn_M) @ self.J_l.T) @ (self.J_l @ np.linalg.inv(self.dyn_M) @ (tauAll - self.dyn_Non) + self.dJ_l @ self.dq)
        self.FRest = -pinv(self.J_r @ np.linalg.inv(self.dyn_M) @ self.J_r.T) @ (self.J_r @ np.linalg.inv(self.dyn_M) @ (tauAll - self.dyn_Non) + self.dJ_r @ self.dq)

        if self.motionState == MotionState.Walk2Stand:
            self.enableNextStep = False
            if self.touchDown:
                self.motionState = MotionState.Stand

        if self.motionState == MotionState.Stand:
            self.dPhi = 0
            self.phi = 0
            self.isIni = False
            self.legState = LegState.RSt
            self.enableNextStep = False
        elif self.motionState == MotionState.Walk:
            self.enableNextStep = True
            self.dPhi = 1.0 / self.tSwing * self.dt
        elif self.motionState == MotionState.Walk2Stand:
            self.dPhi = 1.0 / self.tSwing * self.dt

        self.phi += self.dPhi
        if self.enableNextStep:
            self.touchDown = False

        if not self.isIni:
            self.isIni = True
            if self.legState == LegState.LSt:
                self.swingStartPos_W = self.fe_r_pos_W
                self.stanceStartPos_W = self.fe_l_pos_W
            else:
                self.swingStartPos_W = self.fe_l_pos_W
                self.stanceStartPos_W = self.fe_r_pos_W

        if self.legState == LegState.LSt and self.FRest[2] >= self.FzThrehold and self.phi >= 0.6:
            if self.enableNextStep:
                self.legState = LegState.RSt
                self.swingStartPos_W = self.fe_l_pos_W
                self.stanceStartPos_W = self.fe_r_pos_W
                self.phi = 0
        elif self.legState == LegState.RSt and self.FLest[2] >= self.FzThrehold and self.phi >= 0.6:
            if self.enableNextStep:
                self.legState = LegState.LSt
                self.swingStartPos_W = self.fe_r_pos_W
                self.stanceStartPos_W = self.fe_l_pos_W
                self.phi = 0

        if not self.enableNextStep:
            if self.legState == LegState.LSt and self.FRest[2] >= 200:
                touch_down = True
            if self.legState == LegState.RSt and self.FLest[2] >= 200:
                touch_down = True

        if self.phi >= 1:
            self.phi = 1

        if self.legState == LegState.LSt:
            self.posHip_W = self.hip_r_pos_W
            self.pos_ST_W = self.fe_l_pos_W
            self.theta0 = -3.1415 * 0.5
            self.legStateNext = LegState.RSt
        else:
            self.posHip_W = self.hip_l_pos_W
            self.posST_W = self.fe_r_pos_W
            self.theta0 = 3.1415 * 0.5
            self.legStateNext = LegState.LSt
            
            